package com.zf.agentdefine;

import javassist.ClassPool;
import javassist.CtClass;
import javassist.CtMethod;
import javassist.LoaderClassPath;

import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.security.ProtectionDomain;
import java.util.Objects;
import java.util.StringJoiner;

/**
 * 定义一个 ClassFileTransformer， 给每个方法添加 打印执行时间的打印
 */
public class TimerClassFileTransformer implements ClassFileTransformer {
    public static final String TRANSFORM_CLASS = "com.zf.agentuse.Student";


    /**
     * 在类第一次加载、重定义时调用该方法，可以对类文件进行转换，实现动态代码修改。
     * 如果类不需要被转换，则返回null即可
     *
     * @param loader              将要被重新定义（修改）的类的类加载器，如果是 bootstrapClassLoader则是 null
     * @param className           将要被重新定义（修改）的类的类名，在 Java 虚拟机规范中定义的完全限定类和接口名称的内部形式的类名，如："java/util/List"
     * @param classBeingRedefined 如果在类加载的时候触发的该方法，则为 null;
     *                            如果是重新定义（修改）的类，则是将要被重新定义（修改）的类的 java.lang.Class 对象;
     * @param protectionDomain    将要被重新定义（修改）的类的保护信息
     * @param classfileBuffer     将要被重新定义（修改）的类的字节码字节数组
     */
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
        String fullClassName = className.replace("/", ".");
        System.out.printf("ClassFileTransformer#transform()方法执行了，被转换的类是：%s\n", fullClassName);

        // 不是 指定的类，则不做转换，返回null即可
        if (!Objects.equals(fullClassName, TRANSFORM_CLASS)) {
            return null;
        }

        try {
            // 使用javassist修改类的字节码
            ClassPool classPool = ClassPool.getDefault();

            // 如果是 使用agent修改正在运行中的JVM进程，必须将要修改的类的类加载器所在的类路径加入到 ClassPool 中，否则报错 javassist.NotFoundException: com.zf.agentuse.Student
            // 如果是 启动JVM进程时通过 -javaagent: 参数传递了 agent，则可以不用
            classPool.appendClassPath(new LoaderClassPath(loader));

            // 以javassist格式 获取要修改的类
            CtClass targetCtClass = classPool.get(fullClassName);
            CtClass longCtClass = classPool.get(long.class.getName());

            // 如果类已经执行过 targetCtClass.toBytecode() 或 targetCtClass.writeFile()方法，则类不能再被修改了
            if (targetCtClass.isFrozen()) {
                // 返回上次修改的结果
                return targetCtClass.toBytecode();
            }

            // 获取定义的所有方法，继承的方法不会获取
            CtMethod[] declaredMethods = targetCtClass.getDeclaredMethods();
            for (CtMethod ctMethod : declaredMethods) {
                String name = ctMethod.getName();
                String startTime = String.format("%s$$$StartTime", name);
                String endTime = String.format("%s$$$EndTime", name);

                // 给方法声明两个变量
                ctMethod.addLocalVariable(startTime, longCtClass);
                ctMethod.addLocalVariable(endTime, longCtClass);

                // 在方法的最开始增加语句
                ctMethod.insertBefore(
                        new StringJoiner("\n")
                                .add(String.format("%s = System.currentTimeMillis();", startTime))
                                .toString()

                );

                // 在方法的最后且返回之前增加语句
                ctMethod.insertAfter(
                        new StringJoiner("\n")
                                .add(String.format("%s = System.currentTimeMillis();", endTime))
                                .add(String.format("System.out.println(\"方法 %s 执行耗时：\" + (%s - %s) + \"ms\");", name, endTime, startTime))
                                .toString()
                );
            }

            // 将修改后的字节码输出到文件，输出后就不能在修改了类了；指定目录即可，会自动在目录下输出报名格式的目录，如 D:\\com\\zf\\agentuse\\Student.class
            // targetCtClass.writeFile("D:\\");

            // 返回修改后的字节码，执行该方法后就不能再修改类了
            return targetCtClass.toBytecode();
        } catch (Exception e) {
            e.printStackTrace();
            throw new RuntimeException(e);
        }
    }
}